"""
分布式优化问题定义模块

定义了分布式优化问题的数据结构和相关操作
"""

import numpy as np
from typing import List, Callable, Optional, Tuple, Dict, Any
import cvxpy as cp


class LocalFunction:
    """局部函数类，表示节点i的局部函数"""
    
    def __init__(self, node_id: int, dimension: int):
        """
        初始化局部函数
        
        参数:
            node_id: 节点ID
            dimension: 决策变量维度
        """
        self.node_id = node_id
        self.dimension = dimension
        self._func = None
        self._grad = None
        self._prox = None
        
    def set_function(self, func: Callable, grad: Optional[Callable] = None, 
                    prox: Optional[Callable] = None):
        """设置函数及其梯度和近端算子"""
        self._func = func
        self._grad = grad
        self._prox = prox
        
    def __call__(self, x: np.ndarray) -> float:
        """计算函数值"""
        if self._func is None:
            raise ValueError("函数未定义")
        return self._func(x)
        
    def gradient(self, x: np.ndarray) -> np.ndarray:
        """计算梯度"""
        if self._grad is None:
            raise NotImplementedError("梯度函数未定义")
        return self._grad(x)
        
    def proximal(self, x: np.ndarray, t: float) -> np.ndarray:
        """计算近端算子"""
        if self._prox is None:
            raise NotImplementedError("近端算子未定义")
        return self._prox(x, t)


class DistributedProblem:
    """
    分布式优化问题类
    
    minimize    Σ_{i∈V} f_i(x_i)
    subject to  Σ_{i∈V} g_i(x_i) ≤ 0_m
               Σ_{i∈V} h_i(x_i) = 0_p
               x_i ∈ X_i, ∀i ∈ V
    """
    
    def __init__(self, num_nodes: int, dimensions: List[int] = None):
        """
        初始化分布式优化问题
        
        参数:
            num_nodes: 节点数量
            dimensions: 各节点决策变量维度列表
        """
        self.num_nodes = num_nodes
        if dimensions is None:
            dimensions = [1] * num_nodes  # 默认每个节点维度为1
        self.dimensions = dimensions
        
        # 初始化局部函数
        self.objectives = {}  # 目标函数 f_i
        self.inequalities = {}  # 不等式约束 g_i
        self.equalities = {}  # 等式约束 h_i
        self.local_sets = {}  # 局部约束集 X_i
        
        # 约束维度
        self.m = 0  # 不等式约束数量
        self.p = 0  # 等式约束数量
        
        # 初始化节点
        for i in range(num_nodes):
            self.objectives[i] = LocalFunction(i, dimensions[i])
            self.inequalities[i] = []
            self.equalities[i] = []
            self.local_sets[i] = None
            
    def set_objective(self, node_id: int, func: Callable, 
                     grad: Optional[Callable] = None,
                     prox: Optional[Callable] = None):
        """设置节点的目标函数"""
        self.objectives[node_id].set_function(func, grad, prox)
        
    def add_inequality_constraint(self, node_id: int, func: Callable,
                                 grad: Optional[Callable] = None):
        """添加不等式约束"""
        constraint = LocalFunction(node_id, self.dimensions[node_id])
        constraint.set_function(func, grad)
        self.inequalities[node_id].append(constraint)
        
    def add_equality_constraint(self, node_id: int, func: Callable,
                               grad: Optional[Callable] = None):
        """添加等式约束（必须是仿射函数）"""
        constraint = LocalFunction(node_id, self.dimensions[node_id])
        constraint.set_function(func, grad)
        self.equalities[node_id].append(constraint)
        
    def set_local_constraint_set(self, node_id: int, constraint_set: Dict[str, Any]):
        """
        设置局部约束集
        
        参数:
            node_id: 节点ID
            constraint_set: 约束集描述，例如:
                {'type': 'box', 'lower': -1, 'upper': 1}
                {'type': 'ball', 'center': [0, 0], 'radius': 1}
                {'type': 'simplex', 'dimension': 3}
        """
        self.local_sets[node_id] = constraint_set
        
    def finalize(self):
        """完成问题定义，计算约束维度"""
        # 计算不等式约束总数
        max_ineq = 0
        for i in range(self.num_nodes):
            if len(self.inequalities[i]) > max_ineq:
                max_ineq = len(self.inequalities[i])
        self.m = max_ineq
        
        # 计算等式约束总数
        max_eq = 0
        for i in range(self.num_nodes):
            if len(self.equalities[i]) > max_eq:
                max_eq = len(self.equalities[i])
        self.p = max_eq
        
    def evaluate_objective(self, x: Dict[int, np.ndarray]) -> float:
        """计算总目标函数值"""
        total = 0.0
        for i in range(self.num_nodes):
            if i in x:
                total += self.objectives[i](x[i])
        return total
        
    def evaluate_constraints(self, x: Dict[int, np.ndarray]) -> Tuple[np.ndarray, np.ndarray]:
        """
        计算约束值
        
        返回:
            g: 不等式约束值 (m,)
            h: 等式约束值 (p,)
        """
        g = np.zeros(self.m)
        h = np.zeros(self.p)
        
        # 计算不等式约束
        for j in range(self.m):
            g_j = 0.0
            for i in range(self.num_nodes):
                if j < len(self.inequalities[i]) and i in x:
                    g_j += self.inequalities[i][j](x[i])
            g[j] = g_j
            
        # 计算等式约束
        for j in range(self.p):
            h_j = 0.0
            for i in range(self.num_nodes):
                if j < len(self.equalities[i]) and i in x:
                    h_j += self.equalities[i][j](x[i])
            h[j] = h_j
            
        return g, h
        
    def check_feasibility(self, x: Dict[int, np.ndarray], tol: float = 1e-6) -> bool:
        """检查解的可行性"""
        g, h = self.evaluate_constraints(x)
        
        # 检查不等式约束
        if np.any(g > tol):
            return False
            
        # 检查等式约束
        if np.any(np.abs(h) > tol):
            return False
            
        # 检查局部约束集
        for i in range(self.num_nodes):
            if i in x and self.local_sets[i] is not None:
                if not self._check_local_constraint(x[i], self.local_sets[i]):
                    return False
                    
        return True
        
    def _check_local_constraint(self, x: np.ndarray, constraint_set: Dict[str, Any]) -> bool:
        """检查局部约束集"""
        set_type = constraint_set['type']
        
        if set_type == 'box':
            lower = np.array(constraint_set['lower'])
            upper = np.array(constraint_set['upper'])
            return np.all(x >= lower) and np.all(x <= upper)
            
        elif set_type == 'ball':
            center = np.array(constraint_set['center'])
            radius = constraint_set['radius']
            return np.linalg.norm(x - center) <= radius
            
        elif set_type == 'simplex':
            return np.all(x >= 0) and np.sum(x) <= 1
            
        else:
            raise ValueError(f"未知的约束集类型: {set_type}")
            
    def project_to_local_set(self, x: np.ndarray, node_id: int) -> np.ndarray:
        """投影到局部约束集"""
        if self.local_sets[node_id] is None:
            return x
            
        constraint_set = self.local_sets[node_id]
        set_type = constraint_set['type']
        
        if set_type == 'box':
            lower = np.array(constraint_set['lower'])
            upper = np.array(constraint_set['upper'])
            return np.clip(x, lower, upper)
            
        elif set_type == 'ball':
            center = np.array(constraint_set['center'])
            radius = constraint_set['radius']
            diff = x - center
            norm = np.linalg.norm(diff)
            if norm <= radius:
                return x
            else:
                return center + radius * diff / norm
                
        elif set_type == 'simplex':
            # 投影到单纯形的算法
            return self._project_to_simplex(x)
            
        else:
            raise ValueError(f"未知的约束集类型: {set_type}")
            
    def _project_to_simplex(self, x: np.ndarray) -> np.ndarray:
        """投影到单纯形"""
        n = len(x)
        if np.all(x >= 0) and np.sum(x) <= 1:
            return x
            
        # 使用排序算法投影到单纯形
        u = np.sort(x)[::-1]
        cssv = np.cumsum(u)
        rho = np.where(u > (cssv - 1) / np.arange(1, n + 1))[0][-1]
        theta = (cssv[rho] - 1) / (rho + 1)
        return np.maximum(x - theta, 0) 